#!/usr/bin/env python3
"""
Extract and deduplicate all send.bin files from a tlsfuzzer save directory.

Usage:
  python tools/extract_send_bins.py /path/to/save /path/to/output_dir

Behavior:
  - Recursively searches for files named exactly 'send.bin' under the input
    directory.
  - Deduplicates by SHA-256 content hash.
  - Copies unique files into the output directory named '<sha256>.bin'.
  - Writes a manifest CSV mapping sources to hashes and destination paths.
"""


#  python3 ./tools/extract_send_bins.py ./save/ ./seeds



from __future__ import annotations

import argparse
import csv
import hashlib
import os
from pathlib import Path
import shutil
import sys
from datetime import datetime


def compute_sha256(file_path: Path, chunk_size: int = 1024 * 1024) -> str:
    sha256 = hashlib.sha256()
    with file_path.open('rb') as f:
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
            sha256.update(chunk)
    return sha256.hexdigest()


def find_send_bins(root: Path) -> list[Path]:
    return [p for p in root.rglob('send.bin') if p.is_file()]


def ensure_output_dir(out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)


def write_manifest(manifest_path: Path, rows: list[dict]) -> None:
    fieldnames = [
        'sha256', 'size_bytes', 'source_path', 'copied', 'dest_path'
    ]
    with manifest_path.open('w', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description='Extract and deduplicate send.bin files.'
    )
    parser.add_argument(
        'save_dir', type=Path,
        help='Path to the tlsfuzzer save directory (will be scanned recursively)'
    )
    parser.add_argument(
        'out_dir', type=Path,
        help='Output directory to place unique send.bin files and manifest.csv'
    )
    parser.add_argument(
        '--dry-run', action='store_true',
        help='Only report what would be done without copying files'
    )
    return parser.parse_args(argv)


def main(argv: list[str]) -> int:
    args = parse_args(argv)
    save_dir: Path = args.save_dir
    out_dir: Path = args.out_dir
    dry_run: bool = args.dry_run

    if not save_dir.exists() or not save_dir.is_dir():
        print(f"[!] save_dir does not exist or is not a directory: {save_dir}", file=sys.stderr)
        return 2

    if not dry_run:
        ensure_output_dir(out_dir)

    manifest_rows: list[dict] = []
    seen_hashes: set[str] = set()

    candidates = find_send_bins(save_dir)
    print(f"[*] Found {len(candidates)} send.bin candidate(s) under {save_dir}")

    copied_count = 0
    skipped_dups = 0
    for idx, src in enumerate(sorted(candidates)):
        try:
            size_bytes = src.stat().st_size
        except OSError:
            size_bytes = -1
        sha = compute_sha256(src)

        dest_path = out_dir / f"{sha}.bin"
        is_new = sha not in seen_hashes and not dest_path.exists()

        if is_new and not dry_run:
            # Copy file content to hash-named file
            shutil.copy2(src, dest_path)
            copied = True
            copied_count += 1
            seen_hashes.add(sha)
        else:
            copied = False
            if sha not in seen_hashes:
                seen_hashes.add(sha)
            skipped_dups += 1

        manifest_rows.append({
            'sha256': sha,
            'size_bytes': size_bytes,
            'source_path': str(src.resolve()),
            'copied': str(copied).lower(),
            'dest_path': str(dest_path.resolve()) if not dry_run else '',
        })

        if (idx + 1) % 50 == 0:
            print(f"    processed {idx + 1}/{len(candidates)} ...")

    manifest_path = out_dir / 'manifest.csv'
    if not dry_run:
        write_manifest(manifest_path, manifest_rows)

    print(f"[*] Unique files copied: {copied_count}")
    print(f"[*] Duplicates skipped: {skipped_dups}")
    if not dry_run:
        print(f"[*] Output directory: {out_dir}")
        print(f"[*] Manifest: {manifest_path}")

    return 0

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))


